import talib as ta
import pandas as pd
import backtrader as bt
from collections import namedtuple
from pyLibs.Gadfly import StockStore, StockIndicator, BackTraderIndicator as bti


class ThisStrategy(bt.Strategy):
    params = (
        ('resample_tactics', True),
        ('resample_trend', True),
        ('volume_s_period', 8),
        ('volume_l_period', 100),
        ('ema_s_period', 34),
        ('ema_m_period', 55),
        ('ema_l_period', 144),
        ('boll_period', 21),
        ('boll_devfactor', 2.0),
        ('macd_me1', 12),
        ('macd_me2', 26),
        ('macd_sig', 9),
        ('fastk_period', 9),
        ('slowk_period', 3),
        ('slowd_period',3),
        ('sar_period', 4),
        ('sar_af', 0.02),
        ('sar_afmax', 0.2),
        ('atr_period', 14),
        ('atr_multiple', 1.5),
        ('increase_period', 5),
    )

    def __init__(self):
        p = self.params
        self.indicators = list()
        self.condition = list()
        self.recognition = list()
        self.multiple = list()
        Trade = namedtuple('Trade', ['orderSign', 'buyPrice', 'sellPrice', 'lossPrice'])
        Multiple = namedtuple('Multiple', ['code', 'tacticsData', 'tacticsCycle', 'trendData', 'trendCycle'])
        self.watch = pd.DataFrame(columns=['datetime', 'open', 'close', 'high', 'low', 'ema_s', 'ema_m', 'ema_l',
                                           'boll_mid', 'atr_loss', 'morphology', 'platform_high', 'platform_low',
                                           'vol_release', 'vol_multiplier', 'initial_loss', 'moving_lost',
                                           'break_price', 'status'])
        self.trade = Trade(False, None, None, None)
        tdx = StockStore.TdxStore()
        for i, data in enumerate(self.datas):
            if data._timeframe == bt.TimeFrame.Days:
                if self.params.resample_tactics:
                    tacticsData = self._indicator(tdx.get_stock_data(data._name[2:], '1W').data, '1W')
                else:
                    tacticsData = None
                if self.params.resample_trend:
                    trendData = self._indicator(tdx.get_stock_data(data._name[2:], '1M').data, '1M')
                else:
                    trendData = None
                tacticsCycle = '1W'
                trendCycle = '1M'
            elif data._timeframe == bt.TimeFrame.Minutes:
                if data._compression == 5:
                    if self.params.resample_tactics:
                        tacticsData = self._indicator(tdx.lc5Resample(data._dataname, '15T'), '15T')
                    else:
                        tacticsData = None
                    if self.params.resample_trend:
                        trendData = self._indicator(tdx.lc5Resample(data._dataname, '60T'), '60T')
                    else:
                        trendData = None
                    tacticsCycle = '15T'
                    trendCycle = '60T'
                elif data._compression == 15:
                    if self.params.resample_tactics:
                        tacticsData = self._indicator(tdx.lc5Resample(data._dataname, '60T'), '60T')
                    else:
                        tacticsData = None
                    if self.params.resample_trend:
                        trendData = self._indicator(tdx.get_stock_data(data._name[2:], '1D').data, '1D')
                    else:
                        trendData = None
                    tacticsCycle = '60T'
                    trendCycle = '1D'
                elif data._compression == 60:
                    if self.params.resample_tactics:
                        tacticsData = self._indicator(tdx.get_stock_data(data._name[2:], '1D').data, '1D')
                    else:
                        tacticsData = None
                    if self.params.resample_trend:
                        trendData = self._indicator(tdx.get_stock_data(data._name[2:], '1W').data, '1W')
                    else:
                        trendData = None
                    tacticsCycle = '1D'
                    trendCycle = '1W'
            else:
                tacticsData = None
                trendData = None
                tacticsCycle = None
                trendCycle = None
            self.multiple.append(Multiple(data._name, tacticsData, tacticsCycle, trendData, trendCycle))
            bollinger = bt.ind.BollingerBands(data, period=p.boll_period, devfactor=p.boll_devfactor)
            macd = bt.ind.MACDHisto(data, period_me1=p.macd_me1, period_me2=p.macd_me2, period_signal=p.macd_sig)
            atr = bt.ind.AverageTrueRange(data, period=p.atr_period)
            vol_ma_s = bt.ind.SMA(data.volume, period=p.volume_s_period)
            vol_ma_l = bt.ind.SMA(data.volume, period=p.volume_l_period)
            ema_s = bt.ind.ExponentialMovingAverage(data, period=p.ema_s_period)
            ema_m = bt.ind.ExponentialMovingAverage(data, period=p.ema_m_period)
            ema_l = bt.ind.ExponentialMovingAverage(data, period=p.ema_l_period)
            mid_sma = bt.ind.SMA(bollinger.mid, period=p.increase_period)
            ema_s_sma = bt.ind.SMA(ema_s, period=p.increase_period)
            ema_m_sma = bt.indi.SMA(ema_m, period=p.increase_period)
            self.indicators.append({
                'vol_ma_s': vol_ma_s,
                'vol_ma_l': vol_ma_l,
                'ema_s': ema_s,
                'ema_m': ema_m,
                'ema_l': ema_l,
                'mid': bollinger.mid,
                'upper': bollinger.top,
                'lower': bollinger.bot,
                'bb': (data.close - bollinger.mid) / (bollinger.top - bollinger.bot),       # 布林极限指标
                'width': (bollinger.top - bollinger.bot) / bollinger.mid,                   # 布林带宽度指标
                'dif': macd.macd,
                'dea': macd.signal,
                'macd': macd.histo,
                'sar': bt.ind.ParabolicSAR(data, period=p.sar_period, af=p.sar_af, afmax=p.sar_afmax),
                'atr': atr,
                'atr_loss': bt.ind.Highest((data.low - atr * p.atr_multiple), period=p.atr_period),
                'max_high': bt.ind.Highest(data.high, period=p.atr_period),           # 区间最高价
                'min_low': bt.ind.Lowest(data.low, period=p.atr_period),              # 区间最低价
            })
            self.condition.append({
                'is_vol_release': bt.And(data.volume > vol_ma_s, vol_ma_s > vol_ma_l),                  # 是否放量
                'is_vol_multiplier': bt.And(data.volume > vol_ma_l * 2, data.volume > data.volume[-1]),   # 是否倍量
                'is_short_long_trend': bt.And(ema_s > ema_m, ema_s >= ema_s_sma),   # 是否短期多头趋势
                'is_middle_long_trend': bt.And(ema_m > ema_l, ema_m >= ema_m_sma),  # 是否中期多头趋势
                'is_mid_increase': bollinger.mid > mid_sma,                         # 布林中轨是否上升
                'is_ema_s_increase': ema_s >= ema_s_sma,                            # 中期 EMA 均线是否上升
                'is_ema_m_increase': ema_m >= ema_m_sma,                            # 长期 EMA 均线是否上升
            })
            self.recognition.append({
                'Unique_3_River': bt.talib.CDLUNIQUE3RIVER(data.open, data.high, data.low, data.close),
                'pinbar': bt.talib.CDLHAMMER(data.open, data.high, data.low, data.close)
            })

    def getMultipleData(self, index, x=0):
        pd_data = self.datas[index]._dataname
        if self.datas[index]._timeframe == bt.TimeFrame.Minutes:
            cycle = str(self.datas[index]._compression) + 'T'
        elif self.datas[index]._timeframe == bt.TimeFrame.Days:
            cycle = str(self.datas[index]._compression) + 'D'
        elif self.datas[index]._timeframe == bt.TimeFrame.Weeks:
            cycle = str(self.datas[index]._compression) + 'W'
        else:
            cycle = str(self.datas[index]._compression)
        pd_data = self._indicator(pd_data, cycle)
        datetime = self.datas[index].datetime.datetime(x)
        action = pd_data.query("index == @datetime").tail(1)
        if self.params.resample_tactics and self.params.resample_trend:
            tactics = self.multiple[index].tacticsData.query("index <= @datetime").tail(1)
            trend = self.multiple[index].trendData.query("index <= @datetime").tail(1)
        else:
            if self.params.resample_tactics:
                tactics = self.multiple[index].tacticsData.query("index <= @datetime").tail(1)
                trend = pd.DataFrame()
            elif self.params.resample_trend:
                tactics = pd.DataFrame()
                trend = self.multiple[index].trendData.query("index <= @datetime").tail(1)
            else:
                tactics = pd.DataFrame()
                trend = pd.DataFrame()
        return pd.concat([action, tactics, trend], ignore_index=True)

    def _indicator(self, data, cycle):
        inds = StockIndicator.Indicator(data)
        p = self.params
        return inds.ema(short=p.ema_s_period, middle=p.ema_m_period, long=p.ema_l_period)\
            .bollinger(period=p.boll_period, devfactor=p.boll_devfactor)\
            .volume(short=p.volume_s_period, long=p.volume_l_period)\
            .macd(short=p.macd_me1, long=p.macd_me2, signal=p.macd_sig)\
            .kdj(fastk_period=p.fastk_period, slowk_period=p.slowk_period, slowd_period=p.slowd_period)\
            .atr(period=p.atr_period, multiplier=p.atr_multiple)\
            .get(cycle)

    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.datetime(0)
            print('%s: %s' % (dt.isoformat(sep=" "), txt))

    def printLog(self, log_data, dt=None, doprint=False):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.datetime(0)
            message = '：%s(%s)，关注价位：%.2f，收盘价：%.2f，对比价位：%.2f，区间最低价：%.2f，区间最高价：%.2f，ATR止损价：%.2f'
            print(dt.isoformat(sep=" ") + message % log_data)

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            #self.log('Order Submitted/Accepted')
            return
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(
                    '买入 执行, 交易单价: %.2f, 交易数量: %i，当前账户现金余额: %.2f, 交易金额: %.2f，佣金：%.2f' %
                    (order.executed.price,
                     order.executed.size,
                     self.broker.get_cash(),
                     order.executed.value,
                     order.executed.comm))
                self.trade.buyPrice = order.executed.price
                self.trade.sellPrice = None
                self.trade.orderSign = True
            else:  # Sell
                self.log('卖出 执行, 交易单价: %.2f, 交易数量: %i, 当前账户现金余额: %.2f, 交易金额: %.2f，佣金：%.2f' %
                         (order.executed.price,
                          order.executed.size,
                          self.broker.get_cash(),
                          order.executed.value,
                          order.executed.comm))
                self.trade.buyPrice = None
                self.trade.sellPrice = order.executed.price
                self.trade.orderSign = False
        elif order.status == order.Canceled:
            self.log('Order Canceled')
        elif order.status == order.Margin:
            self.log('Order Margin')
        elif order.status == order.Rejected:
            self.log('Order Rejected')

    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        self.log('\033[31m 交易利润, 总收入: %.2f, 净利润: %.2f \033[0m' % (trade.pnl, trade.pnlcomm))


class BaseStrategy(bt.Strategy):
    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.datetime(0)
            print('%s: %s' % (dt.isoformat(sep=" "), txt))

    def printLog(self, log_data, dt=None, doprint=False):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.datetime(0)
            message = '：%s(%s)，关注价位：%.2f，收盘价：%.2f，对比价位：%.2f，区间最低价：%.2f，区间最高价：%.2f，ATR止损价：%.2f'
            print(dt.isoformat(sep=" ") + message % log_data)

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            # self.log('Order Submitted/Accepted')
            return
        if order.status == order.Completed:
            if order.isbuy():
                self.log(
                    '买入 执行, 交易单价: %.2f, 交易数量: %i，当前账户现金余额: %.2f, 交易金额: %.2f，佣金：%.2f' %
                    (order.executed.price,
                     order.executed.size,
                     self.broker.get_cash(),
                     order.executed.value,
                     order.executed.comm))
                self.trade['buyPrice'] = order.executed.price
                self.trade['sellPrice'] = None
                self.trade['orderSign'] = True
            else:  # Sell
                self.log('卖出 执行, 交易单价: %.2f, 交易数量: %i, 当前账户现金余额: %.2f, 交易金额: %.2f，佣金：%.2f' %
                         (order.executed.price,
                          order.executed.size,
                          self.broker.get_cash(),
                          order.executed.value,
                          order.executed.comm))
                self.trade['platformLow'] = None
                self.trade['lossPrice'] = None
                self.trade['buyPrice'] = None
                self.trade['sellPrice'] = order.executed.price
                self.trade['orderSign'] = False
        elif order.status == order.Canceled:
            if order.isbuy():
                self.log('Buy Order Margin')
            else:
                self.log('Sell Order Margin')
        elif order.status == order.Margin:
            if order.isbuy():
                self.log('Buy Order Margin')
            else:
                self.log('Sell Order Margin')
        elif order.status == order.Rejected:
            if order.isbuy():
                self.log('Buy Order Rejected')
            else:
                self.log('Sell Order Rejected')

    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        self.log('\033[31m 交易利润, 总收入: %.2f, 净利润: %.2f \033[0m' % (trade.pnl, trade.pnlcomm))


class SingleCycleStrategy(BaseStrategy):
    params = (
        ('printlog', True),
        ('volume_period', 100),
        ('ema_s_period', 13),
        ('ema_m_period', 34),
        ('ema_l_period', 55),
        ('boll_period', 20),
        ('boll_multiple', 2.0),
        ('boll_bb_period', 6),
        ('atr_period', 14),
        ('atr_multiple', 1.5),
        ('platform_period', 21),
        ('increase_period', 5),
    )

    columns_e = "datetime platform_upper platform_lower stop_surplus initial_loss " \
                "moving_lost break_price status explain".split()

    columns_c = "t时间日期 p区间上轨 p区间下轨 p前期上轨 l移动止损 l初始止损 l盈亏平衡 s当前状态 e说明解释".split()

    def __init__(self):
        p = self.params
        self.order = None
        self.indicators = list()
        self.condition = list()
        self.watchList = dict()
        self.tradeList = dict()
        for i, data in enumerate(self.datas):
            self.watchList[data._name] = pd.DataFrame(columns=self.columns_c)
            self.tradeList[data._name] = {'orderSign': False, 'buyPrice': None, 'sellPrice': None,
                                          'platformHigh': None, 'platformLow': None, 'lossPrice': None}
            ema_s = bt.ind.EMA(data.close, period=p.ema_s_period)
            ema_m = bt.ind.EMA(data.close, period=p.ema_m_period)
            ema_l = bt.ind.EMA(data.close, period=p.ema_l_period)
            bollinger = bti.BollBand(data.close, maperiod=p.boll_period, multiple=p.boll_multiple)
            boll_bb = bti.BollBB(data.close, boll_period=p.boll_period,
                                 boll_multiple=p.boll_multiple, period=p.boll_bb_period)
            volume_ma = bt.ind.SMA(data.volume, period=p.volume_period)
            mid_sma = bt.ind.SMA(bollinger.middle, period=p.increase_period)
            ema_s_sma = bt.ind.SMA(ema_s, period=p.increase_period)
            ema_m_sma = bt.ind.SMA(ema_m, period=p.increase_period)
            ema_l_sma = bt.ind.SMA(ema_l, period=p.increase_period)
            self.indicators.append({
                'volume_ma': volume_ma,
                'ema_s': ema_s,
                'ema_m': ema_m,
                'ema_l': ema_l,
                'middle': bollinger.middle,
                'upper': bollinger.upper,
                'lower': bollinger.lower,
                'bb': boll_bb.bb,
                'bb_ma': boll_bb.bb_ma,
                'width': bti.BollWidth(data.close, boll_period=p.boll_period, boll_multiple=p.boll_multiple),
                'atr_loss': bti.AtrLoss(data, period=p.atr_period, multiple=p.atr_multiple),
                'platform_upper_high': bt.ind.MaxN(data.high, period=p.platform_period),         # N 个周期内的最高最高价
                'platform_upper_close': bt.ind.MaxN(data.close, period=p.platform_period),       # N 个周期内的最高收盘价
                'platform_lower_low': bt.ind.MinN(data.low, period=p.platform_period),          # N 个周期内的最低最低价
                'platform_lower_close': bt.ind.MinN(data.close, period=p.platform_period),       # N 个周期内的最低收盘价
            })
            self.condition.append({
                # 真阳线放量
                'is_vol_release': bt.And(data.volume > volume_ma * 2, data.ratio > 0, data.close > data.open),
                # 阳线穿越短期均线和布林带中轨
                'kCrossEmaMid': bt.And(bt.ind.CrossUp(data.close, ema_s), bt.ind.CrossUp(data.close, bollinger.middle)),
                'ema_s_cross_up_mid': bt.ind.CrossUp(ema_s, bollinger.middle),  # 短期均线金叉中轨
                # 收盘价金叉短期均线
                'close_cross_ema_s': bt.And(bt.ind.CrossUp(data.close, ema_s), ema_s > bollinger.middle),
                'is_mid_increase': bollinger.middle > mid_sma,  # 布林中轨是否上升
                'is_ema_s_increase': ema_s >= ema_s_sma,  # 短期 EMA 均线是否上升
                'is_ema_m_increase': ema_m >= ema_m_sma,  # 中期 EMA 均线是否上升
                'is_ema_l_increase': ema_l >= ema_l_sma,  # 长期 EMA 均线是否上升
            })


class DoubleCycleStrategy(BaseStrategy):
    params = (
        ('printlog', True),
        ('volume_s_period', 8),
        ('volume_l_period', 100),
        ('ema_s_period', 34),
        ('ema_m_period', 55),
        ('ema_l_period', 144),
        ('boll_period', 21),
        ('boll_factor', 2.0),
        ('macd_me1', 12),
        ('macd_me2', 26),
        ('macd_sig', 9),
        ('fastk_period', 9),
        ('slowk_period', 3),
        ('slowd_period', 3),
        ('sar_period', 4),
        ('sar_af', 0.02),
        ('sar_afmax', 0.2),
        ('atr_period', 14),
        ('atr_multiple', 1.5),
        ('increase_period', 5),
    )

    def __init__(self):
        p = self.params
        self.indicators = list()
        self.condition = list()
        self.recognition = list()
        self.order = None
        self.watchList = dict()
        self.trade = {'orderSign': False, 'buyPrice': None, 'sellPrice': None,
                      'platformHigh': None, 'platformLow': None, 'lossPrice': None}
        for i, data in enumerate(self.datas):
            self.watchList[data._name] = pd.DataFrame(
                columns=['datetime', 'open', 'close', 'high', 'low', 'platform_upper', 'bollinger_upper',
                         'platform_lower', 'stop_surplus', 'initial_loss', 'moving_lost', 'break_price',
                         'status', 'explain'])
            bollinger = bt.ind.BBands(data, period=p.boll_period, devfactor=p.boll_factor)
            macd = bt.ind.MACDHistogram(data, period_me1=p.macd_me1, period_me2=p.macd_me2, period_signal=p.macd_sig)
            atr = bt.ind.ATR(data, period=p.atr_period)
            vol_ma_s = bt.ind.SMA(data.volume, period=p.volume_s_period)
            vol_ma_l = bt.ind.SMA(data.volume, period=p.volume_l_period)
            ema_s = bt.ind.EMA(data, period=p.ema_s_period)
            ema_m = bt.ind.EMA(data, period=p.ema_m_period)
            ema_l = bt.ind.EMA(data, period=p.ema_l_period)
            mid_sma = bt.ind.SMA(bollinger.mid, period=p.increase_period)
            ema_s_sma = bt.ind.SMA(ema_s, period=p.increase_period)
            ema_m_sma = bt.ind.SMA(ema_m, period=p.increase_period)
            self.indicators.append({
                'vol_ma_s': vol_ma_s,
                'vol_ma_l': vol_ma_l,
                'ema_s': ema_s,
                'ema_m': ema_m,
                'ema_l': ema_l,
                'mid': bollinger.mid,
                'upper': bollinger.top,
                'lower': bollinger.bot,
                'bb': (data.close - bollinger.mid) / (bollinger.top - bollinger.bot),  # 布林极限指标
                'width': (bollinger.top - bollinger.bot) / bollinger.mid,  # 布林带宽度指标
                'dif': macd.macd,
                'dea': macd.signal,
                'macd': macd.histo,
                'sar': bt.ind.PSAR(data, period=p.sar_period, af=p.sar_af, afmax=p.sar_afmax),
                'atr': atr,
                'atr_loss': bt.ind.MaxN((data.low - atr * p.atr_multiple), period=p.atr_period),
                # 'platform_upper': btind.MaxN(data.high, period=p.atr_period),  # N 个周期内的最高价
                # 'platform_lower': btind.MinN(data.low, period=p.atr_period),  # N 个周期内的最低价
                # 'stop_surplus': btind.MaxN(data.close, period=p.atr_period),
            })
            morphology_rise_trend = bt.And(data.high_trend == data.platform_upper_trend, bt.Or(
                data.CDLINVERTEDHAMMER > 0, data.CDLSHOOTINGSTAR > 0, data.CDLDARKCLOUDCOVER > 0,
                data.CDLEVENINGSTAR > 0, data.CDLHANGINGMAN > 0, data.CDL3BLACKCROWS > 0
            ))
            morphology_drop_trend = bt.And(data.low_trend == data.platform_lower_trend, bt.Or(
                data.CDLHAMMER > 0, data.CDLINVERTEDHAMMER > 0, data.CDLMORNINGDOJISTAR > 0, data.CDLUNIQUE3RIVER > 0
            ))
            self.condition.append({
                'is_vol_release': bt.And(data.volume > vol_ma_s, vol_ma_s > vol_ma_l),  # 是否放量
                'is_vol_multiplier': bt.And(data.volume > vol_ma_l * 2, data.volume > data.volume[-1]),  # 是否倍量
                'is_vol_multiplier_trend': data.volume_trend > data.volume_ma_trend,  # 是否多头趋势倍量
                'close_cross_ema_s_trend': bt.And(data.low_trend < data.ema_s_trend,
                                                  data.ema_s_trend < data.close_trend),  # 趋势周期收盘金叉EMA_S
                'close_cross_mid_trend': bt.And(data.low_trend < data.middle_trend,
                                                data.middle_trend < data.close_trend),  # 趋势周期收盘金叉MID
                'is_short_long_trend': bt.And(data.ema_s_trend > data.ema_m_trend,
                                              data.middle_trend > data.ema_m_trend),  # 是否短期多头趋势
                'is_middle_long_trend': bt.And(data.ema_s_trend > data.ema_m_trend,
                                               data.ema_m_trend > data.ema_l_trend),  # 是否中期多头趋势
                'morphology_rise_trend': morphology_rise_trend,
                'morphology_drop_trend': morphology_drop_trend,
                'ema_s_cross_up_mid': bt.ind.CrossUp(ema_s, bollinger.mid),  # 短期均线金叉中轨
                'close_cross_ema_s': bt.And(btind.CrossUp(data.close, ema_s), ema_s > bollinger.mid),  # 收盘价金叉短期均线
                'is_mid_increase': bollinger.mid > mid_sma,  # 布林中轨是否上升
                'is_ema_s_increase': ema_s >= ema_s_sma,  # 中期 EMA 均线是否上升
                'is_ema_m_increase': ema_m >= ema_m_sma,  # 长期 EMA 均线是否上升
            })
            self.recognition.append({
                'Unique_3_River': bt.talib.CDLUNIQUE3RIVER(data.open, data.high, data.low, data.close),
                'pinbar': bt.talib.CDLHAMMER(data.open, data.high, data.low, data.close)
            })